/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#include <gtest/gtest.h>
#include "ascend_ir/ascend_ir.h"
#include "common/checker.h"
#include "mmpa/mmpa_api.h"
#include "graph/expression/const_values.h"
#include "graph/utils/graph_utils.h"
#include "graph/ascend_ir/utils/asc_graph_utils.h"

class AscIrRecoveryUtilsUT : public testing::Test {
 protected:
  void SetUp() {
    dlog_setlevel(0, 0, 0);
  }

  void TearDown() {
    dlog_setlevel(0, 3, 0);
  }
};
namespace ge {
namespace {
struct AscNodeInfo {
  std::string name;
  std::string type;
  size_t input_num;
  size_t output_num;
};
NodePtr BuildNode(const ComputeGraphPtr &asc_graph,
                  const AscNodeInfo &node) {
  static int64_t order{0};
  OpDescBuilder op_desc_builder(node.name, node.type);
  for (size_t input_index = 0U; input_index < node.input_num; ++input_index) {
    op_desc_builder.AddInput("x_" + std::to_string(input_index));
  }
  // 当前仅存在单输出的场景
  EXPECT_EQ(node.output_num, 1U);
  op_desc_builder.AddOutput("y");
  auto op_desc = op_desc_builder.Build();
  auto attr_group_of_op = op_desc->GetOrCreateAttrsGroup<AscNodeAttr>();
  EXPECT_NE(attr_group_of_op, nullptr);
  attr_group_of_op->sched.exec_order = order++;
  return asc_graph->AddNode(op_desc);
}

void ConstructAscComputeGraph(AscGraph &graph) {
  Expression s0 = graph.CreateSizeVar("s0");
  // 也可以创建常量，跟var成图的方式不太一样
  Expression s1 = graph.CreateSizeVar(2);

  Axis &s0_axis = graph.CreateAxis("S0", s0);
  Axis &s1_axis = graph.CreateAxis("S1", s1);
  const auto compute_graph = AscGraphUtils::GetComputeGraph(graph);
  auto data0 = BuildNode(compute_graph, {"data0", "Data", 0, 1});
  EXPECT_NE(data0->GetOpDesc(), nullptr);
  data0->GetOpDesc()->GetOrCreateAttrsGroup<AscNodeAttr>()->sched.axis = {s0_axis.id, s1_axis.id};
  auto data0_tensor_desc = data0->GetOpDesc()->MutableOutputDesc(0U);
  data0_tensor_desc->SetDataType(ge::DT_FLOAT16);
  data0_tensor_desc->SetFormat(ge::FORMAT_ND);
  auto data0_tensor_attr_group = data0_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>();
  EXPECT_NE(data0_tensor_attr_group, nullptr);
  data0_tensor_attr_group->axis = {s0_axis.id, s1_axis.id};
  data0_tensor_attr_group->repeats = {s0, s1};
  data0_tensor_attr_group->strides = {s1, sym::kSymbolOne};
  auto data1 = BuildNode(compute_graph, {"data1", "Data", 0, 1});
  data1->GetOpDesc()->GetOrCreateAttrsGroup<AscNodeAttr>()->sched.axis = {s0_axis.id, s1_axis.id};
  auto data1_tensor_desc = data1->GetOpDesc()->MutableOutputDesc(0U);
  data1_tensor_desc->SetDataType(ge::DT_FLOAT16);
  data1_tensor_desc->SetFormat(ge::FORMAT_ND);
  data1_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->axis = {s0_axis.id, s1_axis.id};
  data1_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->repeats = {s0, s1};
  data1_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->strides = {s1, sym::kSymbolOne};
  auto add = BuildNode(compute_graph, {"add", "Add", 2, 1});
  add->GetOpDesc()->GetOrCreateAttrsGroup<AscNodeAttr>()->sched.axis = {s0_axis.id, s1_axis.id};
  auto add_tensor_desc = add->GetOpDesc()->MutableOutputDesc(0U);
  add_tensor_desc->SetDataType(ge::DT_FLOAT16);
  add_tensor_desc->SetFormat(ge::FORMAT_ND);
  add_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->axis = {s0_axis.id, s1_axis.id};
  add_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->repeats = {s0, s1};
  add_tensor_desc->GetOrCreateAttrsGroup<AscTensorAttr>()->strides = {s1, sym::kSymbolOne};
  ge::GraphUtils::AddEdge(data0->GetOutDataAnchor(0), add->GetInDataAnchor(0));
  ge::GraphUtils::AddEdge(data1->GetOutDataAnchor(0), add->GetInDataAnchor(1));
}

Status LoadSo(const char *so_path) {
  GELOGI("Start to dlopen so %s", so_path);
  const uint32_t open_flag = static_cast<uint32_t>(MMPA_RTLD_NOW) | static_cast<uint32_t>(MMPA_RTLD_GLOBAL);
  auto handle = mmDlopen(so_path, static_cast<int32_t>(open_flag));
  GE_ASSERT_NOTNULL(handle, dlerror());
  return SUCCESS;
}
}
TEST_F(AscIrRecoveryUtilsUT, test_recovery_when_convert_case0) {
  AscGraph graph("test");
  ConstructAscComputeGraph(graph);
  auto add_before = AscGraphUtils::GetComputeGraph(graph)->FindNode("add");
  EXPECT_TRUE(add_before != nullptr);
  EXPECT_EQ(add_before->GetOpDesc()->GetIrInputs().size(), 0U);
  EXPECT_EQ(add_before->GetOpDesc()->GetIrOutputs().size(), 0U);
  EXPECT_EQ(LoadSo(std::string(CMAKE_BINARY_DIR).append("/tests/ut/ascend_ir/stub/libascir_stub_builtin_ops.so").c_str()),
            SUCCESS);
  auto graph_after_covert = AscGraph::ConvertComputeGraphToAscGraph(AscGraphUtils::GetComputeGraph(graph));
  auto add = graph_after_covert.FindNode("add");
  EXPECT_TRUE(add != nullptr);
  EXPECT_EQ(add->GetOpDesc()->GetIrInputs().size(), 2U);
  EXPECT_EQ(add->GetOpDesc()->GetIrOutputs().size(), 1U);
}
}